/*
 * Project CelestialBeing: serializable_ptr.hpp
 * Created: 2023/12/15
 * This program uses MulanPSL2. See LICENSE for more.
 */

// An implementation of serializable ptr.
// Serializable ptr can only point to unique objects.

#ifndef CELESTIALBEING_OBJECT_PTR_HPP
#define CELESTIALBEING_OBJECT_PTR_HPP

#include <memory>
#include "base_meta.hpp"
#include "unique_object.hpp"
#include "datatype_id.hpp"

CB_NAMESPACE_BEGIN


template<class T> class TWeakSerializablePtr;
// A serializable shared_ptr for states.
// We don't want to make this class virtual, so no inheritance from Serializable.

// The object id is loaded from the registry upon deserialization, however, the real pointer
// is lazily loaded from the registry upon the FIRST TIME the pointer is USED.
// During other times, the pointer is not really synced with the registry.
// Serialize will trigger a sync and has specific behavior under different conditions.
// SetLoadFrom will reset the lazy load operation, and the pointer will be loaded from the registry again next time
// someone trys to access the underlying pointer.
template<class T>
class TSharedSerializablePtr {
    using CheckTValidType = std::conditional<TIsComplete<T>::value, std::is_base_of<UniqueObject, T>, std::true_type>::type;
public:
    using kValueType = T;
    static_assert(CheckTValidType::value && "Don't use this pointer on any non-state objects. ");

    // Constructors
    inline TSharedSerializablePtr() = default;
    inline TSharedSerializablePtr (std::nullptr_t) : wrapped_(nullptr) {}
    inline TSharedSerializablePtr (const TSharedSerializablePtr & other) = default;
    inline TSharedSerializablePtr (TSharedSerializablePtr && other) noexcept = default;
    inline TSharedSerializablePtr (const std::shared_ptr<T> & ptr) : wrapped_(ptr) {
        SetUniqueObjectIDFromPtr();
    }
    inline TSharedSerializablePtr (std::shared_ptr<T> && ptr) noexcept : wrapped_(std::move(ptr)) {
        SetUniqueObjectIDFromPtr();
    }
    inline explicit TSharedSerializablePtr (UniqueObjectID id) {
        LoadFrom(id);
    }
    // We do not support construction & assignment from raw C pointers.
    // All allocations of UniqueObject should be managed by UniqueObjectRegistry.
//    inline explicit TSharedSerializablePtr (T * ptr) : std::shared_ptr<T>(ptr) {
//        SetUniqueObjectIDFromPtr();
//    }


    // Assignment operators
    inline TSharedSerializablePtr & operator = (std::nullptr_t) {
        wrapped_ = nullptr;
        load_object_id_ = UniqueObjectID::kInvalid;
        lazy_load_scheduled_ = false;
        return *this;
    }
    inline TSharedSerializablePtr & operator = (const TSharedSerializablePtr & other) = default;
    inline TSharedSerializablePtr & operator = (TSharedSerializablePtr && other) noexcept = default;
    inline TSharedSerializablePtr & operator = (const std::shared_ptr<T> & ptr) {
        wrapped_ = ptr;
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TSharedSerializablePtr & operator = (std::shared_ptr<T> && ptr) noexcept {
        wrapped_ = std::move(ptr);
        SetUniqueObjectIDFromPtr();
        return *this;
    }
        // We do not support construction & assignment from raw C pointers for safety reasons.
        // All allocations of UniqueObject should be managed by UniqueObjectRegistry.
//    inline TSharedSerializablePtr & operator = (T * ptr) {
//        std::shared_ptr<T>::operator = (ptr);
//        SetUniqueObjectIDFromPtr();
//        return *this;
//    }

    inline T & operator * () const {
        TryUpdatePointer();
        assert(wrapped_.get());
        return *wrapped_;
    }
    inline operator bool () const {
        TryUpdatePointer();
        return wrapped_.operator bool();
    }
    inline T * operator -> () const {
        TryUpdatePointer();
        return wrapped_.get();
    }
    inline T * Get () const {
        TryUpdatePointer();
        return wrapped_.get();
    }
    inline void Reset () {
        wrapped_.reset();
        load_object_id_ = UniqueObjectID::kInvalid;
        lazy_load_scheduled_ = false;
    }

    std::shared_ptr<T> GetSharedPtr () const {
        TryUpdatePointer();
        return wrapped_;
    }

    // Serialize using the currently present pointer.
    // Update the pointer first if a lazy load is scheduled.
    inline void Serialize (std::ostream & os) const {
        TryUpdatePointer();
        {
            auto referring_object_by_id = UniqueObjectRegistry::Get().GetObject<UniqueObject>(load_object_id_);
            auto referring_object_by_ptr = wrapped_.get();
            // They should always be the same! Unless this pointer is never loaded.
            assert(!referring_object_by_ptr || referring_object_by_id.get() == referring_object_by_ptr);
            load_object_id_ = referring_object_by_id ? referring_object_by_id->GetObjectID() : UniqueObjectID::kInvalid;
        }
        os.write(reinterpret_cast<const char *>(&load_object_id_), sizeof(load_object_id_));
    }
    inline void Deserialize (std::istream & is) {
        is.read(reinterpret_cast<char *>(&load_object_id_), sizeof(load_object_id_));
        LoadFrom(load_object_id_);
    }
    // We banned the SetLoadFrom function. It potentially causes frees on wanted objects due to
    // the lack of synchronization between ID and pointer.
    // inline void SetLoadFrom(UniqueObjectID id) const {

    // Load the pointer from a specific object id.
    // Schedule a lazy load if failed.
    inline bool LoadFrom (UniqueObjectID id) {
        if(id != load_object_id_) wrapped_.reset();
        load_object_id_ = id;
        if(auto ptr = UniqueObjectRegistry::Get().GetObject<T>(load_object_id_)) {
            wrapped_ = ptr;
            lazy_load_scheduled_ = false;
            return true;
        }
        lazy_load_scheduled_ = true;
        return false;
    }
    inline void ScheduleLazyLoad () const {
        lazy_load_scheduled_ = true;
    }


    static constexpr char kTypeNameFormat[] = "TSharedSerializablePtr<{}>";
    // The identity of the pointed type is considered the type/data identity of this type,
    // but the data identity of the pointed type (that is, the internal datastructure of the type)
    // is not considered relevant to the pointer's type/data identity.
    CB_DATATYPE_DECL_INPLACE(FORMAT_1P(kTypeNameFormat, TGetDataTypeMeta<T>::value.name), 0);


    friend class TWeakSerializablePtr<T>;

protected:

    std::shared_ptr<T> & Raw () {
        return wrapped_;
    }
    const std::shared_ptr<T> & Raw () const {
        return wrapped_;
    }

    inline void TryUpdatePointer () const {
        if(lazy_load_scheduled_) {
            wrapped_ = UniqueObjectRegistry::Get().GetObject<T>(load_object_id_);
            assert((load_object_id_ == UniqueObjectID::kInvalid || wrapped_.get()) && "Unique object use before load, or the object is already released.");
            lazy_load_scheduled_ = false;
        }
    }
    inline void SetUniqueObjectIDFromPtr () {
        if(!wrapped_.get()) load_object_id_ = UniqueObjectID::kInvalid;
        else load_object_id_ = wrapped_.get()->GetObjectID();
    }

    mutable bool lazy_load_scheduled_ {true};
    mutable UniqueObjectID load_object_id_ {UniqueObjectID::kInvalid};
    mutable std::shared_ptr<T> wrapped_ {};
};

// A serializable weak_ptr for states.
// We don't want to make this class virtual, so no inheritance from Serializable.

// The object id is loaded from the registry upon deserialization, however, the real pointer
// is lazily loaded from the registry upon the FIRST TIME the pointer is USED.
// During other times the pointer is not really synced with the registry.
// Serialize and Deserialize will trigger a sync.
// SetLoadFrom will reset the lazy load operation, and the pointer will be loaded from the registry again next time
// someone trys to access the underlying pointer.
template<class T>
class TWeakSerializablePtr {
    using CheckTValidType = std::conditional<TIsComplete<T>::value, std::is_base_of<UniqueObject, T>, std::true_type>::type;
public:
    using kValueType = T;
    static_assert(CheckTValidType::value && "Don't use this pointer on any non-state objects. ");

    // Constructors
    inline TWeakSerializablePtr() = default;
    inline TWeakSerializablePtr (std::nullptr_t) : wrapped_(nullptr) {}
    inline TWeakSerializablePtr (const TWeakSerializablePtr & other) = default;
    inline TWeakSerializablePtr (TWeakSerializablePtr && other) noexcept = default;
    inline TWeakSerializablePtr (const TSharedSerializablePtr<T> & other) : wrapped_(other.Raw()) {
        SetUniqueObjectIDFromPtr();
    }
    inline TWeakSerializablePtr (TSharedSerializablePtr<T> && other) noexcept : wrapped_(std::move(other.Raw())) {
        SetUniqueObjectIDFromPtr();
    }
    inline TWeakSerializablePtr (const std::shared_ptr<T> & ptr) : wrapped_(ptr) {
        SetUniqueObjectIDFromPtr();
    }
    inline TWeakSerializablePtr (std::shared_ptr<T> && ptr) noexcept : wrapped_(std::move(ptr)) {
        SetUniqueObjectIDFromPtr();
    }
    inline TWeakSerializablePtr (const std::weak_ptr<T> & ptr) : wrapped_(ptr) {
        SetUniqueObjectIDFromPtr();
    }
    inline TWeakSerializablePtr (std::weak_ptr<T> && ptr) noexcept : wrapped_(std::move(ptr)) {
        SetUniqueObjectIDFromPtr();
    }
    inline explicit TWeakSerializablePtr (UniqueObjectID id) {
        LoadFrom(id);
    }
    // We do not support construction & assignment from raw C pointers.
    // All allocations of UniqueObject should be managed by UniqueObjectRegistry.
//    inline explicit TWeakSerializablePtr (T * ptr) : std::weak_ptr<T>(ptr) {
//        SetUniqueObjectIDFromPtr();
//    }

    // Assignment operators
    inline TWeakSerializablePtr & operator = (std::nullptr_t) {
        wrapped_ = nullptr;
        load_object_id_ = UniqueObjectID::kInvalid;
        lazy_load_scheduled_ = false;
        return *this;
    }
    inline TWeakSerializablePtr & operator = (const TWeakSerializablePtr & other) = default;
    inline TWeakSerializablePtr & operator = (TWeakSerializablePtr && other) noexcept = default;
    inline TWeakSerializablePtr & operator = (const std::shared_ptr<T> & ptr) {
        wrapped_ = ptr;
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TWeakSerializablePtr & operator = (std::shared_ptr<T> && ptr) noexcept {
        wrapped_ = std::move(ptr);
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TWeakSerializablePtr & operator = (const TSharedSerializablePtr<T> & other) {
        std::weak_ptr<T>::operator = (other);
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TWeakSerializablePtr & operator = (TSharedSerializablePtr<T> && other) noexcept {
        std::weak_ptr<T>::operator = (std::move(other));
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TWeakSerializablePtr & operator = (const std::weak_ptr<T> & ptr) {
        std::weak_ptr<T>::operator = (ptr);
        SetUniqueObjectIDFromPtr();
        return *this;
    }
    inline TWeakSerializablePtr & operator = (std::weak_ptr<T> && ptr) noexcept {
        std::weak_ptr<T>::operator = (std::move(ptr));
        SetUniqueObjectIDFromPtr();
        return *this;
    }
        // We do not support construction & assignment from raw C pointers for safety reasons.
        // All allocations of UniqueObject should be managed by UniqueObjectRegistry.
//    inline TWeakSerializablePtr & operator = (T * ptr) {
//        std::weak_ptr<T>::operator = (ptr);
//        SetUniqueObjectIDFromPtr();
//        return *this;
//    }

    inline TSharedSerializablePtr<T> Lock () const {
        TryUpdatePointer();
        return TSharedSerializablePtr<T>(wrapped_.lock());
    }

    inline bool Expired () const {
        TryUpdatePointer();
        return wrapped_.expired();
    }

    inline void Reset () {
        wrapped_.reset();
        load_object_id_ = UniqueObjectID::kInvalid;
        lazy_load_scheduled_ = false;
    }

    inline std::weak_ptr<T> GetWeakPtr () const {
        TryUpdatePointer();
        return wrapped_;
    }

    // Serialize using the currently present pointer.
    // Update the pointer first if a lazy load is scheduled.
    inline void Serialize (std::ostream & os) const {
        TryUpdatePointer();
        {
            auto referring_object_by_id = UniqueObjectRegistry::Get().GetObject<UniqueObject>(load_object_id_);
            auto referring_object_by_ptr = wrapped_.lock.get();
            // They should always be the same! Unless this pointer is never loaded.
            assert(!referring_object_by_ptr || referring_object_by_id.get() == referring_object_by_ptr);
            load_object_id_ = referring_object_by_id ? referring_object_by_id->GetObjectID() : UniqueObjectID::kInvalid;
        }
        os.write(reinterpret_cast<const char *>(&load_object_id_), sizeof(load_object_id_));
    }
    inline void Deserialize (std::istream & is) {
        is.read(reinterpret_cast<char *>(&load_object_id_), sizeof(load_object_id_));
        LoadFrom(load_object_id_);
    }
    // We banned the SetLoadFrom function. It potentially causes frees on wanted objects due to
    // the lack of synchronization between ID and pointer.
    // inline void SetLoadFrom(UniqueObjectID id) const {...}

    // Load the pointer from a specific object id.
    // Schedule a lazy load if failed.
    inline bool LoadFrom (UniqueObjectID id) {
        if(id != load_object_id_) wrapped_ = nullptr;
        load_object_id_ = id;
        if(auto ptr = UniqueObjectRegistry::Get().GetObject<T>(load_object_id_)) {
            wrapped_ = ptr;
            lazy_load_scheduled_ = false;
            return true;
        }
        lazy_load_scheduled_ = true;
        return false;
    }
    inline void ScheduleLazyLoad () const {
        lazy_load_scheduled_ = true;
    }
    static constexpr char kTypeNameFormat[] = "TWeakSerializablePtr<{}>";
    // The identity of the pointed type is considered the type/data identity of this type,
    // but the data identity of the pointed type (that is, the internal datastructure of the type)
    // is not considered relevant to the pointer's type/data identity.
    CB_DATATYPE_DECL_INPLACE(FORMAT_1P(kTypeNameFormat, TGetDataTypeMeta<T>::value.name), 0);

protected:
    std::weak_ptr<T> & Raw () {
        return wrapped_;
    }
    const std::weak_ptr<T> & Raw () const {
        return wrapped_;
    }
    inline void TryUpdatePointer () const {
        if(lazy_load_scheduled_) {
            auto ptr = UniqueObjectRegistry::Get().GetObject<T>(load_object_id_);
            wrapped_ = ptr;
            assert((load_object_id_ == UniqueObjectID::kInvalid || ptr) && "Unique object use before load.");
            lazy_load_scheduled_ = false;
        }
    }
    inline void SetUniqueObjectIDFromPtr () {
        if(auto ptr = wrapped_.lock()) load_object_id_ = ptr->GetObjectID();
        else load_object_id_ = UniqueObjectID::kInvalid;
    }
    mutable std::weak_ptr<T> wrapped_;

    mutable bool lazy_load_scheduled_ {true};
    mutable UniqueObjectID load_object_id_ {UniqueObjectID::kInvalid};
};

template<typename T> struct TIsSerializablePtr { static constexpr bool value = false; };
template<typename T> struct TIsSerializablePtr<TSharedSerializablePtr<T>> { static constexpr bool value = true; };
template<typename T> struct TIsSerializablePtr<TWeakSerializablePtr<T>> { static constexpr bool value = true; };

CB_NAMESPACE_END

#endif //CELESTIALBEING_OBJECT_PTR_HPP
